#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
实时脑电信号可视化器

本文件用于与eeg_data_logger.py对接，实时显示脑电信号数据。
专门针对频域信号（Delta, Theta, Alpha, Beta, Gamma）进行实时可视化。

使用方法：
1. 确保eeg_data_logger.py正在运行并记录数据
2. 运行此脚本进行实时可视化
3. 支持多种显示模式和数据源

原理：
1. 继承EEGDataLogger类，重写回调函数获取实时数据
2. 使用matplotlib.animation实现图表实时刷新
3. 维护滑动窗口数据缓存，优化显示性能
4. 支持多子图显示不同频段的信号变化
"""

import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'eeg'))

import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.dates import DateFormatter
import numpy as np
from datetime import datetime, timedelta
from collections import deque
import threading
import time
import pandas as pd

try:
    from eeg_data_logger import EEGDataLogger
except ImportError:
    print("警告: 无法导入eeg_data_logger，将使用独立模式")
    EEGDataLogger = None

class RealTimeEEGVisualizer:
    """实时脑电信号可视化器类"""
    
    def __init__(self, port="COM5", baudrate=57600, window_size=30, data_source="live"):
        """
        初始化实时脑电信号可视化器
        
        Args:
            port (str): 串口号，默认COM5
            baudrate (int): 波特率，默认57600
            window_size (int): 显示窗口大小（秒），默认30秒
            data_source (str): 数据源类型 - "live"(实时), "csv"(CSV文件), "simulate"(模拟)
        """
        self.port = port
        self.baudrate = baudrate
        self.window_size = window_size
        self.data_source = data_source
        
        # 数据缓存队列
        self.max_points = window_size * 10  # 假设每秒10个数据点
        self.time_data = deque(maxlen=self.max_points)
        self.attention_data = deque(maxlen=self.max_points)
        self.meditation_data = deque(maxlen=self.max_points)
        self.delta_data = deque(maxlen=self.max_points)
        self.theta_data = deque(maxlen=self.max_points)
        self.low_alpha_data = deque(maxlen=self.max_points)
        self.high_alpha_data = deque(maxlen=self.max_points)
        self.low_beta_data = deque(maxlen=self.max_points)
        self.high_beta_data = deque(maxlen=self.max_points)
        self.low_gamma_data = deque(maxlen=self.max_points)
        self.mid_gamma_data = deque(maxlen=self.max_points)
        self.signal_quality_data = deque(maxlen=self.max_points)
        
        # 当前数据值
        self.current_data = {
            'attention': 0,
            'meditation': 0,
            'delta': 0,
            'theta': 0,
            'lowAlpha': 0,
            'highAlpha': 0,
            'lowBeta': 0,
            'highBeta': 0,
            'lowGamma': 0,
            'midGamma': 0,
            'poorSignal': 200
        }
        
        # 图表相关
        self.fig = None
        self.axes = None
        self.lines = {}
        self.ani = None
        
        # 线程控制
        self.data_lock = threading.Lock()
        self.is_running = False
        self.neuropy = None
        
        # CSV文件路径（如果使用CSV数据源）
        self.csv_file = None
        self.csv_index = 0
    
    def update_data_from_dict(self, data_dict):
        """从数据字典更新当前数据"""
        self.current_data.update(data_dict)
        self._add_data_point()
    
    def _add_data_point(self):
        """添加数据点到缓存队列"""
        timestamp = datetime.now()
        
        with self.data_lock:
            self.time_data.append(timestamp)
            self.attention_data.append(self.current_data['attention'])
            self.meditation_data.append(self.current_data['meditation'])
            self.delta_data.append(self.current_data['delta'])
            self.theta_data.append(self.current_data['theta'])
            self.low_alpha_data.append(self.current_data['lowAlpha'])
            self.high_alpha_data.append(self.current_data['highAlpha'])
            self.low_beta_data.append(self.current_data['lowBeta'])
            self.high_beta_data.append(self.current_data['highBeta'])
            self.low_gamma_data.append(self.current_data['lowGamma'])
            self.mid_gamma_data.append(self.current_data['midGamma'])
            self.signal_quality_data.append(200 - self.current_data['poorSignal'])  # 转换为信号质量
    
    def setup_plots(self):
        """设置图表布局"""
        # 设置中文字体
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False
        
        # 创建图表 - 3x2布局
        self.fig, self.axes = plt.subplots(3, 2, figsize=(16, 12))
        self.fig.suptitle('实时脑电信号可视化器', fontsize=16, fontweight='bold', color='darkblue')
        
        # 子图1: 注意力和冥想状态
        ax1 = self.axes[0, 0]
        ax1.set_title('注意力 & 冥想状态', fontweight='bold')
        ax1.set_ylabel('指数值 (0-100)')
        ax1.set_ylim(0, 100)
        ax1.grid(True, alpha=0.3)
        self.lines['attention'], = ax1.plot([], [], 'r-', label='注意力', linewidth=2)
        self.lines['meditation'], = ax1.plot([], [], 'b-', label='冥想', linewidth=2)
        ax1.legend()
        ax1.xaxis.set_major_formatter(DateFormatter('%H:%M:%S'))
        
        # 子图2: Gamma波（重点显示）
        ax2 = self.axes[0, 1]
        ax2.set_title('Gamma波实时监测', fontweight='bold', color='red')
        ax2.set_ylabel('功率值')
        ax2.grid(True, alpha=0.3)
        self.lines['low_gamma'], = ax2.plot([], [], 'g-', label='低Gamma', linewidth=3)
        self.lines['mid_gamma'], = ax2.plot([], [], 'm-', label='中Gamma', linewidth=3)
        self.lines['total_gamma'], = ax2.plot([], [], 'r-', label='总Gamma', linewidth=2, alpha=0.7)
        ax2.legend()
        ax2.xaxis.set_major_formatter(DateFormatter('%H:%M:%S'))
        
        # 子图3: Alpha波段
        ax3 = self.axes[1, 0]
        ax3.set_title('Alpha波段', fontweight='bold')
        ax3.set_ylabel('功率值')
        ax3.grid(True, alpha=0.3)
        self.lines['low_alpha'], = ax3.plot([], [], 'orange', label='低Alpha', linewidth=2)
        self.lines['high_alpha'], = ax3.plot([], [], 'gold', label='高Alpha', linewidth=2)
        ax3.legend()
        ax3.xaxis.set_major_formatter(DateFormatter('%H:%M:%S'))
        
        # 子图4: Beta波段
        ax4 = self.axes[1, 1]
        ax4.set_title('Beta波段', fontweight='bold')
        ax4.set_ylabel('功率值')
        ax4.grid(True, alpha=0.3)
        self.lines['low_beta'], = ax4.plot([], [], 'cyan', label='低Beta', linewidth=2)
        self.lines['high_beta'], = ax4.plot([], [], 'navy', label='高Beta', linewidth=2)
        ax4.legend()
        ax4.xaxis.set_major_formatter(DateFormatter('%H:%M:%S'))
        
        # 子图5: Delta & Theta波段
        ax5 = self.axes[2, 0]
        ax5.set_title('Delta & Theta波段', fontweight='bold')
        ax5.set_ylabel('功率值')
        ax5.grid(True, alpha=0.3)
        self.lines['delta'], = ax5.plot([], [], 'purple', label='Delta', linewidth=2)
        self.lines['theta'], = ax5.plot([], [], 'brown', label='Theta', linewidth=2)
        ax5.legend()
        ax5.xaxis.set_major_formatter(DateFormatter('%H:%M:%S'))
        
        # 子图6: 信号质量和实时数值
        ax6 = self.axes[2, 1]
        ax6.set_title('信号质量 & 实时数值', fontweight='bold')
        ax6.set_ylabel('信号质量 (0-200)')
        ax6.set_ylim(0, 200)
        ax6.grid(True, alpha=0.3)
        self.lines['signal_quality'], = ax6.plot([], [], 'black', label='信号质量', linewidth=2)
        ax6.legend()
        ax6.xaxis.set_major_formatter(DateFormatter('%H:%M:%S'))
        
        # 添加实时数值显示文本
        self.text_display = ax6.text(0.02, 0.98, '', transform=ax6.transAxes, 
                                   fontsize=8, verticalalignment='top',
                                   bbox=dict(boxstyle='round', facecolor='lightgray', alpha=0.8))
        
        plt.tight_layout()
        plt.subplots_adjust(top=0.95)
    
    def update_plots(self, frame):
        """更新图表数据"""
        with self.data_lock:
            if len(self.time_data) == 0:
                return list(self.lines.values())
            
            # 获取当前数据
            times = list(self.time_data)
            attention = list(self.attention_data)
            meditation = list(self.meditation_data)
            delta = list(self.delta_data)
            theta = list(self.theta_data)
            low_alpha = list(self.low_alpha_data)
            high_alpha = list(self.high_alpha_data)
            low_beta = list(self.low_beta_data)
            high_beta = list(self.high_beta_data)
            low_gamma = list(self.low_gamma_data)
            mid_gamma = list(self.mid_gamma_data)
            signal_quality = list(self.signal_quality_data)
        
        # 计算总Gamma
        total_gamma = [lg + mg for lg, mg in zip(low_gamma, mid_gamma)]
        
        # 更新所有线条数据
        self.lines['attention'].set_data(times, attention)
        self.lines['meditation'].set_data(times, meditation)
        self.lines['low_gamma'].set_data(times, low_gamma)
        self.lines['mid_gamma'].set_data(times, mid_gamma)
        self.lines['total_gamma'].set_data(times, total_gamma)
        self.lines['low_alpha'].set_data(times, low_alpha)
        self.lines['high_alpha'].set_data(times, high_alpha)
        self.lines['low_beta'].set_data(times, low_beta)
        self.lines['high_beta'].set_data(times, high_beta)
        self.lines['delta'].set_data(times, delta)
        self.lines['theta'].set_data(times, theta)
        self.lines['signal_quality'].set_data(times, signal_quality)
        
        # 动态调整x轴范围
        if len(times) > 0:
            current_time = times[-1]
            start_time = current_time - timedelta(seconds=self.window_size)
            
            for ax in self.axes.flat:
                ax.set_xlim(start_time, current_time)
        
        # 动态调整y轴范围（除了固定范围的子图）
        if len(times) > 10:
            # Gamma波子图
            if total_gamma:
                max_gamma = max(total_gamma[-50:]) if len(total_gamma) >= 50 else max(total_gamma)
                if max_gamma > 0:
                    self.axes[0, 1].set_ylim(0, max_gamma * 1.1)
            
            # Alpha波子图
            alpha_values = low_alpha[-50:] + high_alpha[-50:] if len(low_alpha) >= 50 else low_alpha + high_alpha
            if alpha_values:
                max_alpha = max(alpha_values)
                if max_alpha > 0:
                    self.axes[1, 0].set_ylim(0, max_alpha * 1.1)
            
            # Beta波子图
            beta_values = low_beta[-50:] + high_beta[-50:] if len(low_beta) >= 50 else low_beta + high_beta
            if beta_values:
                max_beta = max(beta_values)
                if max_beta > 0:
                    self.axes[1, 1].set_ylim(0, max_beta * 1.1)
            
            # Delta & Theta子图
            dt_values = delta[-50:] + theta[-50:] if len(delta) >= 50 else delta + theta
            if dt_values:
                max_dt = max(dt_values)
                if max_dt > 0:
                    self.axes[2, 0].set_ylim(0, max_dt * 1.1)
        
        # 更新实时数值显示
        if len(times) > 0:
            current_values = {
                'attention': attention[-1] if attention else 0,
                'meditation': meditation[-1] if meditation else 0,
                'low_gamma': low_gamma[-1] if low_gamma else 0,
                'mid_gamma': mid_gamma[-1] if mid_gamma else 0,
                'signal_quality': signal_quality[-1] if signal_quality else 0
            }
            
            text_content = f"""实时数值:
注意力: {current_values['attention']:.0f}
冥想: {current_values['meditation']:.0f}
低Gamma: {current_values['low_gamma']:.0f}
中Gamma: {current_values['mid_gamma']:.0f}
信号质量: {current_values['signal_quality']:.0f}"""
            
            self.text_display.set_text(text_content)
        
        return list(self.lines.values())
    
    def start_simulation(self):
        """启动模拟数据模式"""
        def simulate_data():
            while self.is_running:
                # 生成模拟的脑电波数据
                t = time.time()
                
                # 模拟各种脑电波信号
                self.current_data = {
                    'attention': max(0, min(100, 50 + 30 * np.sin(t * 0.05) + np.random.normal(0, 5))),
                    'meditation': max(0, min(100, 60 + 25 * np.sin(t * 0.03 + 1) + np.random.normal(0, 5))),
                    'delta': max(0, 2000 + 800 * np.sin(t * 0.02) + np.random.normal(0, 200)),
                    'theta': max(0, 1500 + 600 * np.sin(t * 0.03) + np.random.normal(0, 150)),
                    'lowAlpha': max(0, 1200 + 500 * np.sin(t * 0.04) + np.random.normal(0, 120)),
                    'highAlpha': max(0, 1000 + 400 * np.sin(t * 0.05) + np.random.normal(0, 100)),
                    'lowBeta': max(0, 800 + 300 * np.sin(t * 0.06) + np.random.normal(0, 80)),
                    'highBeta': max(0, 600 + 250 * np.sin(t * 0.07) + np.random.normal(0, 60)),
                    'lowGamma': max(0, 1000 + 500 * np.sin(t * 0.1) + np.random.normal(0, 100)),
                    'midGamma': max(0, 800 + 400 * np.sin(t * 0.15 + 1) + np.random.normal(0, 80)),
                    'poorSignal': max(0, min(200, 50 + 30 * np.sin(t * 0.02) + np.random.normal(0, 10)))
                }
                
                self._add_data_point()
                time.sleep(0.1)  # 每100ms更新一次
        
        # 启动模拟数据线程
        self.data_thread = threading.Thread(target=simulate_data)
        self.data_thread.daemon = True
        self.data_thread.start()
        
        print("模拟数据模式已启动")
    
    def start_csv_reading(self, csv_file):
        """启动CSV文件读取模式"""
        self.csv_file = csv_file
        
        def read_csv_data():
            try:
                df = pd.read_csv(csv_file)
                print(f"加载CSV文件: {csv_file}, 共{len(df)}行数据")
                
                while self.is_running and self.csv_index < len(df):
                    row = df.iloc[self.csv_index]
                    
                    # 更新当前数据
                    self.current_data = {
                        'attention': row.get('attention', 0),
                        'meditation': row.get('meditation', 0),
                        'delta': row.get('delta', 0),
                        'theta': row.get('theta', 0),
                        'lowAlpha': row.get('lowAlpha', 0),
                        'highAlpha': row.get('highAlpha', 0),
                        'lowBeta': row.get('lowBeta', 0),
                        'highBeta': row.get('highBeta', 0),
                        'lowGamma': row.get('lowGamma', 0),
                        'midGamma': row.get('midGamma', 0),
                        'poorSignal': row.get('poorSignal', 0)
                    }
                    
                    self._add_data_point()
                    self.csv_index += 1
                    
                    time.sleep(0.1)  # 控制播放速度
                
                if self.csv_index >= len(df):
                    print("CSV数据播放完毕，循环播放")
                    self.csv_index = 0  # 循环播放
                    
            except Exception as e:
                print(f"读取CSV文件错误: {e}")
                self.start_simulation()  # fallback到模拟模式
        
        # 启动CSV读取线程
        self.data_thread = threading.Thread(target=read_csv_data)
        self.data_thread.daemon = True
        self.data_thread.start()
        
        print(f"CSV数据读取模式已启动: {csv_file}")
    
    def start_live_data(self):
        """启动实时数据采集（与eeg_data_logger.py对接）"""
        try:
            if EEGDataLogger is None:
                raise ImportError("无法导入EEGDataLogger")
            
            # 这里可以实现与eeg_data_logger.py的对接逻辑
            # 例如读取其生成的实时CSV文件或通过其他方式获取数据
            
            print("实时数据采集模式启动失败，切换到模拟模式")
            self.start_simulation()
            
        except Exception as e:
            print(f"实时数据采集启动失败: {e}")
            print("切换到模拟数据模式")
            self.start_simulation()
    
    def start_visualization(self):
        """开始可视化"""
        try:
            # 设置图表
            self.setup_plots()
            
            # 启动数据采集
            self.is_running = True
            
            if self.data_source == "simulate":
                self.start_simulation()
            elif self.data_source == "csv":
                csv_file = input("请输入CSV文件路径 (默认: ../clean_data/eeg_data_cleaned.csv): ").strip()
                if not csv_file:
                    csv_file = "../clean_data/eeg_data_cleaned.csv"
                self.start_csv_reading(csv_file)
            else:
                self.start_live_data()
            
            # 启动动画
            self.ani = animation.FuncAnimation(
                self.fig, self.update_plots, interval=100,
                blit=False, cache_frame_data=False
            )
            
            print("=== 实时脑电信号可视化已启动 ===")
            print("关闭图表窗口或按Ctrl+C停止可视化")
            
            # 显示图表
            plt.show()
            
        except KeyboardInterrupt:
            print("\n用户中断，停止可视化...")
            self.stop_visualization()
        except Exception as e:
            print(f"可视化错误: {e}")
            self.stop_visualization()
    
    def stop_visualization(self):
        """停止可视化"""
        self.is_running = False
        if self.ani:
            self.ani.event_source.stop()
        plt.close('all')
        print("实时脑电信号可视化已停止")

def main():
    """主函数"""
    print("=== 实时脑电信号可视化器 ===")
    print("支持多种数据源的脑电信号实时可视化")
    print()
    
    # 询问用户选择数据源
    print("选择数据源:")
    print("1. 实时数据 (与eeg_data_logger.py对接)")
    print("2. CSV文件 (播放已记录的数据)")
    print("3. 模拟数据 (用于测试)")
    
    choice = input("请选择 (1-3): ").strip()
    
    data_source_map = {
        '1': 'live',
        '2': 'csv', 
        '3': 'simulate'
    }
    
    data_source = data_source_map.get(choice, 'simulate')
    
    print(f"选择的数据源: {data_source}")
    
    # 创建可视化器实例
    visualizer = RealTimeEEGVisualizer(
        port="COM5",
        baudrate=57600,
        window_size=30,  # 显示最近30秒的数据
        data_source=data_source
    )
    
    # 启动可视化
    visualizer.start_visualization()

if __name__ == "__main__":
    main()